/*
 * Copyright 2022 The Open Islands Authors. All Rights Reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.openislands.oi.scheduling;

import org.openislands.oi.config.MyNodeInfo;
import org.openislands.oi.constant.JobState;
import org.openislands.oi.constant.SchedulingAction;
import org.openislands.oi.constant.TaskState;
import org.openislands.oi.pojo.qo.*;
import org.openislands.oi.pojo.vo.OfNodeTaskReportVO;
import org.openislands.oi.rpc.Generic;
import org.openislands.oi.pojo.dto.OfNodeTaskDTO;
import org.openislands.oi.pojo.dto.CommonResult;
import org.openislands.oi.pojo.dto.JobDTO;
import org.openislands.oi.pojo.dto.TaskDTO;
import org.openislands.oi.pojo.vo.OfNodeTaskVO;
import org.openislands.oi.pojo.vo.CommonResponse;
import org.openislands.oi.rpc.RPCExecutor;
import org.openislands.oi.util.ValueCopyUtils;
import org.openislands.oi.worker.http.TaskController;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import javax.annotation.Resource;

@Component
public class TaskCoordinationProtocol {
    private static final Logger log = LoggerFactory.getLogger(TaskCoordinationProtocol.class);
    @Resource
    private RPCExecutor rpcExecutor;

    @Autowired
    private MyNodeInfo myNodeInfo;

    public CoordinationResult<Boolean> createJob(JobDTO jobDTO) {
        CoordinationResult<Boolean> result = rpcExecutor.simple(jobDTO.getNodes(), new Generic<TaskController, CommonResponse<Boolean>>(
                params -> params.createJob(ValueCopyUtils.copyProperties(jobDTO, new JobCreateQO()))) {
        });
        actionLog(result, SchedulingAction.CREATE_JOB);
        return result;
    }

    public CoordinationResult<Boolean> deleteJob(JobDTO jobDTO) {
        CoordinationResult<Boolean> result = rpcExecutor.simple(jobDTO.getNodes(), new Generic<TaskController, CommonResponse<Boolean>>(
                params -> params.deleteJob(jobDTO.getJobId())) {
        });
        actionLog(result, SchedulingAction.DELETE_JOB);
        return result;
    }

    public CoordinationResult<Boolean> syncJob(JobDTO jobDTO) {
        CoordinationResult<Boolean> result = rpcExecutor.simple(jobDTO.getNodes(), new Generic<TaskController, CommonResponse<Boolean>>(
                params -> params.updateJob(jobDTO.getJobId(), ValueCopyUtils.copyProperties(jobDTO, new JobUpdateQO()))) {
        });
        actionLog(result, SchedulingAction.SYNC_JOB);
        return result;
    }

    public CoordinationResult<Boolean> syncJobState(JobDTO jobDTO, JobState jobState) {
        CoordinationResult<Boolean> result = rpcExecutor.simple(jobDTO.getNodes(), new Generic<TaskController, CommonResponse<Boolean>>(
                params -> params.updateJobState(jobDTO.getJobId(), jobState)) {
        });
        actionLog(result, SchedulingAction.SYNC_JOB_STATE);
        return result;
    }

    public CoordinationResult<Boolean> stopJob(JobDTO jobDTO, JobState jobState) {
        CoordinationResult<Boolean> result = rpcExecutor.simple(jobDTO.getNodes(), new Generic<TaskController, CommonResponse<Boolean>>(
                params -> params.stopJob(jobDTO.getJobId(), jobState)) {
        });
        actionLog(result, SchedulingAction.STOP_JOB);
        return result;
    }

    public CoordinationResult<Boolean> createTask(TaskDTO taskDTO) {
        CoordinationResult<Boolean> result = rpcExecutor.simple(taskDTO.getNodes(), new Generic<TaskController, CommonResponse<Boolean>>(
                params -> params.createTask(taskDTO.getJobId(), ValueCopyUtils.copyProperties(taskDTO, new TaskCreateQO()))) {
        });
        actionLog(result, SchedulingAction.CREATE_TASK);
        return result;
    }

    public CoordinationResult<Boolean> syncTask(TaskDTO taskDTO) {
        CoordinationResult<Boolean> result = rpcExecutor.simple(taskDTO.getNodes(), new Generic<TaskController, CommonResponse<Boolean>>(
                params -> params.updateTask(taskDTO.getJobId(), taskDTO.getTaskName(), taskDTO.getTaskVersion(), ValueCopyUtils.copyProperties(taskDTO, new TaskUpdateQO()))) {
        });
        actionLog(result, SchedulingAction.SYNC_TASK);
        return result;
    }

    public CoordinationResult<Boolean> syncTaskState(TaskDTO taskDTO, TaskState toState) {
        CoordinationResult<Boolean> result = rpcExecutor.simple(taskDTO.getNodes(), new Generic<TaskController, CommonResponse<Boolean>>(
                params -> params.updateTaskState(taskDTO.getJobId(), taskDTO.getTaskName(), taskDTO.getTaskVersion(), toState)) {
        });
        actionLog(result, SchedulingAction.SYNC_TASK_STATE);
        return result;
    }

    public CoordinationResult<OfNodeTaskDTO> startTask(TaskDTO taskDTO, TaskState toState) {
        // todo: add task execution env to protocol, and then proxy route to the worker by node id and execution env
        CoordinationResult<OfNodeTaskDTO> result = new CoordinationResult<>();
        taskDTO.getNodes().forEach((role, nodes) -> {
            nodes.forEach(nodeId -> {
                CommonResult<OfNodeTaskVO> oneResult = rpcExecutor.simple(nodeId, new Generic<TaskController, CommonResponse<OfNodeTaskVO>>(
                        params -> params.startTask(taskDTO.getJobId(), taskDTO.getTaskName(), taskDTO.getTaskVersion(), role, toState)) {

                });
                CommonResult<OfNodeTaskDTO> transformResult = CommonResult.copy(oneResult, ValueCopyUtils.copyProperties(oneResult.getData(), new OfNodeTaskDTO()));
                result.addOneOfResult(role, nodeId, transformResult);
                result.calculateCoordinationResult();
            });
        });
        actionLog(result, SchedulingAction.SUBMIT_TASK);
        return result;
    }

    public CoordinationResult<OfNodeTaskDTO> queryTaskOfNodes(TaskDTO taskDTO) {
        CoordinationResult<OfNodeTaskDTO> result = new CoordinationResult<>();
        taskDTO.getNodes().forEach((role, nodes) -> {
            nodes.forEach(nodeId -> {
                CommonResult<OfNodeTaskReportVO> oneResult = rpcExecutor.simple(nodeId, new Generic<TaskController, CommonResponse<OfNodeTaskReportVO>>(
                        params -> params.queryTaskOfNode(taskDTO.getJobId(), taskDTO.getTaskName(), taskDTO.getTaskVersion(), role)) {
                });
                CommonResult<OfNodeTaskDTO> transformResult = CommonResult.copy(oneResult, ValueCopyUtils.copyProperties(oneResult.getData(), new OfNodeTaskDTO()));
                result.addOneOfResult(role, nodeId, transformResult);
                result.calculateCoordinationResult();
            });
        });
        actionLog(result, SchedulingAction.COLLECT_TASK);
        return result;
    }

    public CoordinationResult<Boolean> activateTask(TaskDTO taskDTO, TaskState toState) {
        CoordinationResult<Boolean> result = new CoordinationResult<>();
        taskDTO.getNodes().forEach((role, nodes) -> {
            nodes.forEach(nodeId -> {
                CommonResult<Boolean> oneResult = rpcExecutor.simple(nodeId,
                        new Generic<TaskController, CommonResponse<Boolean>>(
                                params -> params.activateTask(taskDTO.getJobId(), taskDTO.getTaskName(), taskDTO.getTaskVersion(), role, toState)) {
                        }
                );
                result.addOneOfResult(role, nodeId, oneResult);
                result.calculateCoordinationResult();
            });
        });
        actionLog(result, SchedulingAction.ACTIVATE_TASK);
        return result;
    }

    public CoordinationResult<Boolean> stopTask(TaskDTO taskDTO, TaskState toState) {
        CoordinationResult<Boolean> result = new CoordinationResult<>();
        taskDTO.getNodes().forEach((role, nodes) -> {
            nodes.forEach(nodeId -> {
                CommonResult<Boolean> oneResult = rpcExecutor.simple(nodeId, new Generic<TaskController, CommonResponse<Boolean>>(
                        params -> params.stopTask(taskDTO.getJobId(), taskDTO.getTaskName(), taskDTO.getTaskVersion(), role, toState)) {
                });
                result.addOneOfResult(role, nodeId, oneResult);
                result.calculateCoordinationResult();
            });
        });
        actionLog(result, SchedulingAction.STOP_TASK);
        return result;
    }

    public CommonResult<Boolean> reportTaskToLeader(OfNodeTaskDTO ofNodeTaskDTO) {
        CommonResult<Boolean> result = rpcExecutor.simple(ofNodeTaskDTO.getLeaderNode(), new Generic<TaskController, CommonResponse<Boolean>>(
                params -> params.reportTask(ofNodeTaskDTO.getJobId(), ofNodeTaskDTO.getTaskName(), ofNodeTaskDTO.getTaskVersion(),
                        ofNodeTaskDTO.getRole(), ofNodeTaskDTO.getNodeId(), ValueCopyUtils.copyProperties(ofNodeTaskDTO, new OfNodeTaskReportQO()))) {
        });
        return result;
    }

    private <E> void actionLog(CoordinationResult<E> result, SchedulingAction action){
        if (!result.isSuccess()){
            log.error("{} failed: {}", action, result.getMessage());
        }
    }
}
